#!/usr/bin/env python3
"""
NLP提示模块

将文本描述转换为坐标点提示
根据表述定位指定心动周期中的帧区间中的可能坐标集合
"""

import torch
import torch.nn as nn
import numpy as np
from typing import Dict, List, Tuple, Optional, Any, Union
import re

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from .base_prompt import BasePromptGenerator
from .input_prompt import InputPromptHandler
from utils.logger import get_logger


class NLPromptGenerator(BasePromptGenerator):
    """NLP提示生成器
    
    将文本描述转换为坐标点提示
    """
    
    def __init__(self,
                 input_dim: int = 256,
                 prompt_dim: int = 256,
                 device: str = "cuda",
                 use_nlp_model: bool = False):
        """
        初始化NLP提示生成器
        
        Args:
            input_dim: 输入特征维度
            prompt_dim: 提示特征维度
            device: 设备
            use_nlp_model: 是否使用NLP模型（如果False，使用规则解析）
        """
        super().__init__(input_dim, prompt_dim, device)
        
        self.use_nlp_model = use_nlp_model
        self.input_handler = InputPromptHandler(device)
        
        # 如果使用NLP模型，可以在这里初始化
        # 目前先使用规则解析
        if use_nlp_model:
            # TODO: 初始化NLP模型（BERT/BioBERT等）
            pass
        
        # 坐标生成网络（从文本特征生成坐标）
        self.coordinate_generator = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, prompt_dim // 4),
            nn.ReLU(),
            nn.Linear(prompt_dim // 4, 2)  # (x, y)
        )
        
        # 多候选生成器（一个文本可能对应多个坐标）
        self.multi_candidate_generator = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim),
            nn.ReLU(),
            nn.Linear(prompt_dim, 10 * 2)  # 最多10个候选坐标
        )
        
        # 置信度评估器
        self.confidence_estimator = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 1),
            nn.Sigmoid()
        )
        
        self.logger.info(f"NLP提示生成器初始化完成")
    
    def parse_text_description(self, text: str) -> Dict[str, Any]:
        """
        解析文本描述
        
        Args:
            text: 文本描述，如"左心室在舒张末期的心尖位置"
            
        Returns:
            结构化信息字典
        """
        # 简单的规则解析（可以后续扩展为NLP模型）
        info = {
            'anatomical_structure': None,  # 解剖结构
            'cardiac_phase': None,  # 心动周期阶段
            'spatial_location': None,  # 空间位置
            'frame_range': None,  # 帧区间
        }
        
        # 提取心动周期阶段
        if '舒张末期' in text or 'ED' in text or 'end-diastolic' in text.lower():
            info['cardiac_phase'] = 'ED'
        elif '收缩末期' in text or 'ES' in text or 'end-systolic' in text.lower():
            info['cardiac_phase'] = 'ES'
        elif '舒张期' in text or 'diastolic' in text.lower():
            info['cardiac_phase'] = 'diastolic'
        elif '收缩期' in text or 'systolic' in text.lower():
            info['cardiac_phase'] = 'systolic'
        
        # 提取解剖结构
        if '左心室' in text or 'left ventricle' in text.lower() or 'LV' in text:
            info['anatomical_structure'] = 'LV'
        elif '右心室' in text or 'right ventricle' in text.lower() or 'RV' in text:
            info['anatomical_structure'] = 'RV'
        
        # 提取空间位置
        if '心尖' in text or 'apex' in text.lower():
            info['spatial_location'] = 'apex'
        elif '基底' in text or 'base' in text.lower():
            info['spatial_location'] = 'base'
        elif '中部' in text or 'mid' in text.lower():
            info['spatial_location'] = 'mid'
        
        return info
    
    def locate_frame_range(self, 
                          cardiac_phase: Optional[str],
                          video_length: int,
                          frame_indices: Optional[List[int]] = None) -> List[int]:
        """
        定位指定心动周期中的帧区间
        
        Args:
            cardiac_phase: 心动周期阶段 ('ED', 'ES', 'diastolic', 'systolic')
            video_length: 视频总长度
            frame_indices: 可选的帧索引列表（如果有周期检测结果）
            
        Returns:
            帧区间列表
        """
        if frame_indices is not None:
            return frame_indices
        
        # 如果没有提供帧索引，根据阶段返回估计的帧区间
        if cardiac_phase == 'ED':
            # ED通常在视频开始附近
            return list(range(min(5, video_length)))
        elif cardiac_phase == 'ES':
            # ES通常在视频中间附近
            mid = video_length // 2
            return list(range(max(0, mid - 2), min(video_length, mid + 3)))
        elif cardiac_phase == 'diastolic':
            # 舒张期范围较广
            return list(range(video_length))
        elif cardiac_phase == 'systolic':
            # 收缩期范围较广
            return list(range(video_length))
        else:
            # 默认返回所有帧
            return list(range(video_length))
    
    def generate_coordinates(self,
                            text_info: Dict[str, Any],
                            image_shape: Tuple[int, int],
                            num_candidates: int = 5) -> List[Tuple[float, float, float]]:
        """
        生成可能的坐标集合
        
        Args:
            text_info: 解析后的文本信息
            image_shape: 图像形状 (H, W)
            num_candidates: 候选坐标数量
            
        Returns:
            候选坐标列表 [(x, y, confidence), ...]
        """
        H, W = image_shape
        
        # 根据空间位置生成候选坐标
        candidates = []
        
        spatial_location = text_info.get('spatial_location', 'mid')
        
        if spatial_location == 'apex':
            # 心尖位置：通常在下半部分中心
            base_x, base_y = W // 2, int(H * 0.7)
            for i in range(num_candidates):
                x = base_x + np.random.randint(-W//10, W//10)
                y = base_y + np.random.randint(-H//10, H//10)
                x = np.clip(x, 0, W-1)
                y = np.clip(y, 0, H-1)
                confidence = 0.8 - i * 0.1
                candidates.append((float(x), float(y), confidence))
        
        elif spatial_location == 'base':
            # 基底位置：通常在上半部分中心
            base_x, base_y = W // 2, int(H * 0.3)
            for i in range(num_candidates):
                x = base_x + np.random.randint(-W//10, W//10)
                y = base_y + np.random.randint(-H//10, H//10)
                x = np.clip(x, 0, W-1)
                y = np.clip(y, 0, H-1)
                confidence = 0.8 - i * 0.1
                candidates.append((float(x), float(y), confidence))
        
        else:
            # 默认（中部或其他）：图像中心附近
            base_x, base_y = W // 2, H // 2
            for i in range(num_candidates):
                x = base_x + np.random.randint(-W//8, W//8)
                y = base_y + np.random.randint(-H//8, H//8)
                x = np.clip(x, 0, W-1)
                y = np.clip(y, 0, H-1)
                confidence = 0.7 - i * 0.08
                candidates.append((float(x), float(y), confidence))
        
        return candidates
    
    def forward(self,
                image_features: torch.Tensor,
                context: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        前向传播
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            context: 上下文信息（包含text_description等）
            
        Returns:
            提示字典
        """
        # 如果使用NLP模型，这里可以处理文本
        # 目前使用规则解析
        if context and 'text_description' in context:
            text_info = self.parse_text_description(context['text_description'])
            # 生成坐标候选
            image_shape = (image_features.shape[2], image_features.shape[3])
            candidates = self.generate_coordinates(text_info, image_shape)
            
            # 选择置信度最高的候选
            best_candidate = max(candidates, key=lambda x: x[2])
            coords = torch.tensor([[best_candidate[0], best_candidate[1]]], 
                                 dtype=torch.float32, device=self.device)
            
            return self.input_handler.process_point_prompt(coords)
        
        # 如果没有文本描述，返回空提示
        return {}
    
    def generate_prompts(self,
                        text_description: str,
                        image_shape: Tuple[int, int],
                        video_length: int = 1,
                        frame_indices: Optional[List[int]] = None) -> Dict[str, Any]:
        """
        生成提示
        
        Args:
            text_description: 文本描述
            image_shape: 图像形状 (H, W)
            video_length: 视频长度
            frame_indices: 帧索引列表（如果有周期检测结果）
            
        Returns:
            提示字典，包含点提示和相关信息
        """
        # 解析文本
        text_info = self.parse_text_description(text_description)
        
        # 定位帧区间
        frame_range = self.locate_frame_range(
            text_info.get('cardiac_phase'),
            video_length,
            frame_indices
        )
        
        # 生成坐标候选
        candidates = self.generate_coordinates(text_info, image_shape)
        
        # 转换为点提示格式
        # 选择置信度最高的几个候选
        top_candidates = sorted(candidates, key=lambda x: x[2], reverse=True)[:3]
        coords = torch.tensor([[c[0], c[1]] for c in top_candidates],
                             dtype=torch.float32, device=self.device)
        labels = torch.ones(len(top_candidates), dtype=torch.int64, device=self.device)
        
        point_prompt = self.input_handler.process_point_prompt(coords, labels)
        
        return {
            **point_prompt,
            'frame_range': frame_range,
            'cardiac_phase': text_info.get('cardiac_phase'),
            'candidates': candidates,
            'confidence': max(c[2] for c in candidates)
        }
